#!/usr/bin/env python

#from BeautifulSoup import BeautifulSoup
#from Bio.PDB import *
#from optparse import OptionParser

import csv
import os
import os.path
import ctypes
#import MySQLdb
import dbtbl
from dbbuild import *
import sqlite3
#import urllib
# This script will give the protein sequences of PDZ domains that bind to the
# peptides determined by peptide_extract.pl by searching the domino file and
# uniprot database. bind_from and bind_to from domino file and full sequence
# from database.

# 3 things needed to run this script
# 1. peptide1.out file generated by peptide1.py
# 2. domino database mitab file can download from DOMINO database FTP
# 3. Uniprot database (.xml) file in uniprot/ folder

'''
For PDZ <--> peptides
load peptide2 file, using dict. DOMINOID as key, peptide as value
Protein A + proteinB == '-'
link by DOMINO ID
#------------------------------------#
PDZ <--> other proteins
1. find which is the PDZ, A or B
2. find what are the binding sites
     PDZ domain --> more than ONE? overlaps?
     third protein --> binding sites at c-term? almost c-term?
     or somewhere else
3. output interactions

Note:
A-B interactions will have redundences. i.e. overlapping PDZ domains
There will be some cases that B not interacts at the c-term
and very likely cases that interacts near c-terminal.
How to distinguish near c-term and in the middle???19854198

'''


def masking(filename, uni_acc, peptide):
    #lOAD Validation dataset
    validation_table = csv.reader(open(filename))
    for row in validation_table:
        uniprot_acc = row[0]
        vpeptide = row[3]
        if uniprot_acc == uni_acc and peptide == vpeptide[-5:]:
            return True
    return False
    
#-------------------------------------------------------------------##

def import_sqlite(filename, con):
    persons= csv.reader(open("users.csv"))
    con.execute("create table person(firstname, lastname)")
    con.executemany("insert into person(firstname, lastname) values (?, ?)", persons)

def build_newtables(pdzdomains):
    #09/11/2012
    #build the three extra tables protein, pdz2 and taxon
    pdzids = set()
    tbl_protein = set()
    tbl_pdz2 = set()
    tbl_taxon = set()
    for k, v in pdzdomains:
        
        pdzids.add(k)
        
    for acc in pdzids:
        print "Acc:", acc
        genename, seq, organism, taxonid, proteinname = proteininfo(acc)
        
        
        tbl_protein.add((acc, genename[0], proteinname, taxonid, seq))
        pdz2 = pdz2_ext(acc)
        if pdz2 is not False:
            for kk in pdz2:
                tbl_pdz2.add((acc, kk, pdz2[kk][0], pdz2[kk][1]))
        
        #else:
        #    raise ("No PDZ domain found in ", acc)
            
        tbl_taxon.add((taxonid, organism))
        
    return tbl_protein, tbl_taxon, tbl_pdz2


    
    
###================================================================###
###----------------------------------------------------------------###
###================================================================###
#pep = load_peptide('/home/sdai/workspace/PDZ/peptide1.out')

def main():
    print "Loading input files...",
    pep = load_peptide('peptide2.out')
    #domino_table = csv.reader(open(
    #        '/home/sdai/workspace/PDZ/2009-10-22-domino-pdz-binary.mitab26'),
    #        delimiter='\t')
    
    domino_table = csv.reader(open('2009-10-22-domino-pdz-binary.mitab26'),
                              delimiter='\t')

    
    print 'DONE'
    #csvreader.
    uni_proteinID = set()
    proteinIDA = list()
    proteinIDB = list()
    BindingSiteA = list()
    BindingSiteB = list()
    BindingSiteAe = list()
    BindingSiteBe = list()
    BindingSiteAs = list()
    BindingSiteBs = list()
    BindingDomainA = list()
    BindingDomainB =list()
    proteinIDAdb = list()
    proteinIDBdb = list()
    DominoID = list()
    taxidA = list()
    taxidB = list()
    wholeSeq = list()
    
    
    print "Loading Domino Table...",
    for row in domino_table:
        #iiden = row[13] #Interaction identifier(s)
        uni_proteinID.add(row[0].split(':')[1]) #interactor ID A
        
        #uni_proteinIDB.add(row[0].split(':')[1])
        #load binding pocket
        
        try:
            proteinIDA.append(row[0].split(':')[1])
            proteinIDAdb.append(row[0].split(':')[0])
        except (IndexError):
            proteinIDA.append('-')
            proteinIDAdb.append('-')
            
        try:
            proteinIDB.append(row[1].split(':')[1])
            proteinIDBdb.append(row[1].split(':')[0])
        except (IndexError):
            proteinIDB.append('-')
            proteinIDBdb.append('-')
            
            
        try:
            BindingDomainA.append(row[32].split(';')[2])
        except (IndexError):
            BindingDomainA.append('-')
            
        try:
            BindingDomainB.append(row[33].split(';')[2])
        except (IndexError):
            BindingDomainB.append('-')
            
        try: 
            BindingSiteA.append(row[32].split(';')[0])
            
        except (IndexError):
            BindingSiteA.append('-')
    
        try:
            BindingSiteAe.append(int(row[32].split(';')[0].split('-')[1]))
        except (ValueError, IndexError):
            BindingSiteAe.append(0)
            
        try:
            BindingSiteAs.append(int(row[32].split(';')[0].split('-')[0]))
        except (ValueError, IndexError):
            BindingSiteAs.append(0)
            
        try:
            BindingSiteB.append(row[33].split(';')[0])
        except (ValueError, IndexError):
            BindingSiteB.append('-')
    
        try:
            BindingSiteBe.append(int(row[33].split(';')[0].split('-')[1]))
        except (ValueError, IndexError):
            BindingSiteBe.append(0)
    
        try:
            BindingSiteBs.append(int(row[33].split(';')[0].split('-')[0]))
        except (ValueError, IndexError):
            BindingSiteBs.append(0)
            
        try: 
            taxidA.append(row[9].split(':')[1].split('(')[0])
        except:
            taxidA.append('-')
            
        try:
            taxidB.append(row[10].split(':')[1].split('(')[0])
        except:
            taxidB.append('-')
            
                    
        try:
            DominoID.append(row[13].split(':')[1])
        except:
            DominoID.append('-')
    print 'DONE'
    #----------------------#
    
    # 1. load the binding pocket pocket
    
    # 2. extract xx;xx;xx
    
    # 3. extract size from Uniprot (uniprot parser)
            
    # 4. compare size  point to range .....
    
    # 
    #----------------------#
    #interactorAc
    #print len(BindingSiteB)
    #print len(BindingSiteA)
    #print len(proteinIDA)
    print "Processing Binding Sites."
    pdzdomains = dict()
    
    for iter in range(len(BindingSiteA)):
    
        #Protein A + Protein B both PDZ
        #Protein A is, protein B is not
        #Protein A is not, protein B is
        #Protein A + peptide
        
        if (proteinIDAdb[iter] == 'uniprotkb' and
            (re.match('interpro:IPR001478', BindingDomainA[iter]))):
            #Protein A is a PDZ domain
            #print proteinIDA[iter], BindingSiteAs[iter], BindingSiteAe[iter]
            
            (protA, Astart, Aend, Aseq) = uniprot_ext(proteinIDA[iter],
                                                      BindingSiteAs[iter],
                                                      BindingSiteAe[iter])
            if protA:    
                if (proteinIDBdb[iter] == 'uniprotkb'
                    and proteinIDB[iter] != '-'
                    and (re.match('interpro:IPR001478', BindingDomainB[iter]))):
                    #there are only ONE instant of protein A and B are both PDZs.
                    #Q02410 autoinhibition of X11/alpha
                    
                    #try:
                    (protB, Bstart, Bend, Bseq) = uniprot_ext(proteinIDB[iter],
                                                              BindingSiteBs[iter],
                                                              BindingSiteBe[iter])
                    #except(TypeError):
                    if protB:
                        key = (proteinIDA[iter], proteinIDB[iter])
                        if key in pdzdomains:
                            (ks,ke,ksa,ksb) = pdzdomains[key]
                            if (abs(ke-ks) > abs(Aend - Astart)) or (ke == ks):
                                pdzdomains[key] = (Astart, Aend, Aseq, Bseq[-5:])
                        else:
                            pdzdomains[key] = (Astart, Aend, Aseq, Bseq[-5:])
                    else:
                        print 'Protein B:', proteinIDB[iter], ' has been deleted.'
                        
                elif proteinIDB[iter] == '-':
                    #protein A - peptide (ID_B empty)
                    #record protein A in pdzdomains
                    #pep[Domino ID] = (acc, peptide)
                    if DominoID[iter] in pep:
                        key = (pep[DominoID[iter]][0],pep[DominoID[iter]][1][-5:])
                        if key in pdzdomains:
                            (ks,ke,ksa,ksb) = pdzdomains[key]
                            if (abs(ke-ks) > abs(Aend - Astart)) or (ke == ks):
                                pdzdomains[key] = (Astart, Aend, Aseq, key[1])
                        else:
                            pdzdomains[key] = (Astart, Aend, Aseq, key[1])
                            #print proteinIDA[iter] == pep[DominoID[iter]][0]
        
        
                else:
                    #protein A<PDZ, proteinB<not PDZ
                    #record protein A in pdzdomains
                    #print 'A is B not'
                    
                    (protB, Bstart, Bend, Bseq) = uniprot_ext(proteinIDB[iter],
                                                              BindingSiteBs[iter],
                                                              BindingSiteBe[iter])

                    if protB:
                        key = (proteinIDA[iter], proteinIDB[iter])
                        if key in pdzdomains:
                            (ks,ke,ksa,ksb) = pdzdomains[key]
                            if (abs(ke-ks) > abs(Aend - Astart)) or (ke == ks):
                                pdzdomains[key] = (Astart, Aend, Aseq, Bseq[-5:])
                        else:
                            pdzdomains[key] = (Astart, Aend, Aseq, Bseq[-5:])
                    else:
                        print 'Protein B:', proteinIDB[iter], ' has been deleted.'
            else:
                print 'Protein A:', proteinIDA[iter], ' has been deleted.'
                    
        elif (proteinIDAdb[iter] == 'uniprotkb'
              and proteinIDBdb[iter] == 'uniprotikb'
              and (re.match('interpro:IPR001478', BindingDomainB[iter]))):
            #protein A is not PDZ, protein B is PDZ
            (protA, Astart, Aend, Aseq) = uniprot_ext(proteinIDA[iter],
                                                      BindingSiteAs[iter],
                                                      BindingSiteAe[iter])
            (protB, Bstart, Bend, Bseq) = uniprot_ext(proteinIDB[iter],
                                                      BindingSiteBs[iter],
                                                      BindingSiteBe[iter])
            #key will have the pdz domain in front
            key = (proteinIDB[iter], proteinIDA[iter])
            if key in pdzdomains:
                (ks,ke,ksa,ksb) = pdzdomains[key]
                if (abs(ke-ks) > abs(Bend - Bstart)) or (ke == ks):
                    pdzdomains[key] = (Bstart, Bend, seqB, Aseq[-5:])
            else:
                pdzdomains[key] = (Bstart, Bend, seqB, Aseq[-5:])
        else:
            #Non of protein A or B is PDZ domain?
            #do nothing
            pass
            
    
    # download uniprot filen
    #locale database
    #for iter in uni_proteinID:
    #    uniprot_file = urllib.urlretrieve
    #        ('http://www.uniprot.org/uniprot/'+iter+'.xml', iter+'.xml')
    
    #10 fold cross-validation database
    #10 databases: pdzvalid_1 to pdzvalid_10
    #using different mask files

    
    
    print "Building Database..."
    #---------------------------
    #Insert codes for adding three new tables
    tbl_protein, tbl_taxon, tbl_pdz2 = build_newtables(pdzdomains)
    
    #tbl_taxon =
    #tbl_pdz2 = 
    
    #---------------------------
    
    

        
    dbname = "pdzvalid.db"    
    tbl_sdr = set()
    tbl_sdrdom = set()
    tbl_pdzdom = dict()
    tbl_peptide = set()
    tbl_pepdom = set()
    #vfile = 'test_group'+'.pdat'
    for i in pdzdomains:
        #(tbl_sdr, tbl_sdrdom, tbl_pdzdom, tbl_peptide, tbl_pepdom)
        hold = dbtbl.build(
                i[0], i[1], pdzdomains[i][0], pdzdomains[i][1], pdzdomains[i][2],
                pdzdomains[i][3], tbl_sdr, tbl_sdrdom, tbl_pdzdom, tbl_peptide,
                tbl_pepdom)
            #dbtbl.build(protein A, protein B, start position, end position,
            #           sequence, peptide, table holders)
        if hold:
                
            (tbl_sdr, tbl_sdrdom, tbl_pdzdom, tbl_peptide, tbl_pepdom) = hold
    
    #--------
    #look at tbl_pdzdom file --> uniprot
    #print tbl_pdzdom 
    #--------
    #raise

    cwd = os.getcwd()
    # Delete the old table
    '''
    if os.path.isfile(dbname):
        os.remove(dbname)
    '''
    mydb = sqlite3.connect(dbname)
    curs = mydb.cursor()
    sqlFile = 'pdzdb_schema1003.sql'
    
    # Create the tables
    qry = open(sqlFile, 'r').read()
    sqlite3.complete_statement(qry)

    try:
        curs.executescript(qry)
    except Exception as e:
        errorMessage = dbname + ': ' + str(e)
        print errorMessage
        curs.close()
        raise

    '''
    fname = cwd+'/sdr.tbl'
    with open(fname, 'ab') as f:
        for i in tbl_sdr:
            writer = csv.writer(f, delimiter = '\t', lineterminator = '\n')
            writer.writerow(i)
    query = "LOAD DATA LOCAL INFILE %s INTO TABLE sdr;"
    curs.execute(query, (fname))
    '''
    
    fname = cwd+'/sdr_pdz.tbl'

    with open(fname, 'w') as f:
        for i in tbl_sdrdom:
            writer = csv.writer(f, delimiter = '\t', lineterminator = '\n')
            writer.writerow(i)
            
    #query = "LOAD DATA LOCAL INFILE %s INTO TABLE sdr_pdz;"
    query = "INSERT INTO sdr_pdz values (?,?,?,?,?);"
    try:
        curs.executemany(query, tbl_sdrdom)
    except sqlite3.IntegrityError:
        print "tbl_sdrdom: Primary Keys exists."
    
    fname = cwd+'/pdz.tbl'
    with open(fname, 'w') as f:
        for k, v in tbl_pdzdom.iteritems():
            writer = csv.writer(f, delimiter = '\t', lineterminator = '\n')
            writer.writerow((k, ) + v)
            query = "INSERT INTO pdz values (?,?,?,?,?);"
            try:
                curs.execute(query, (k, ) + v)
            except sqlite3.IntegrityError:
                print "tbl_pdz: Primary Key exists."
    
    fname = cwd+'/peptide.tbl'
    with open(fname, 'w') as f:
        for i in tbl_peptide:
            writer = csv.writer(f, delimiter = '\t', lineterminator = '\n')
            writer.writerow(i)
    query = "INSERT INTO peptide values (?,?);"
    try:
        curs.executemany(query, tbl_peptide)
    except sqlite3.IntegrityError:
        print "tbl_peptide: Primary Keys exist."
        
    fname = cwd+'/peptide_pdz'+'.tbl'
    with open(fname, 'w') as f:
        for i in tbl_pepdom:
            writer = csv.writer(f, delimiter = '\t', lineterminator = '\n')
            writer.writerow(i)
    
    query = "INSERT INTO peptide_pdz values (?,?);"
    try:
        curs.executemany(query, tbl_pepdom)
    except sqlite3.IntegrityError:
        print "peptide_pdz: Primary keys exist."
        
    #ADD Three New Tables 09/11/2012
    query = "INSERT INTO protein values (?,?,?,?,?);"
    try:
        curs.executemany(query, tbl_protein)
    except sqlite3.IntegrityError:
        print "protein: Primary keys exist."
        
    query = "INSERT INTO taxon values (?,?);"
    try:
        curs.executemany(query, tbl_taxon)
    except sqlite3.IntegrityError:
        print "taxon: Primary keys exist."
        
    query = "INSERT INTO PDZ2 values (?,?,?,?);"
    try:
        curs.executemany(query, tbl_pdz2)
    except sqlite3.IntegrityError:
        print "PDZ2: Primary keys exist."
    
    mydb.commit()    

#import into database

if __name__ == '__main__':
    print "Starting..."
    main()